# -*- coding: utf-8 -*-
# @File: backend => websocket.py
# @CreateTime: 2025/8/27 13:52
# @Author: EvanZhang
# @Info: WebSocket相关工具

from typing import List
from fastapi import WebSocket
import json

class WebSocketConnectionManager:
    """
    WebSocket连接管理器
    管理所有活动的WebSocket连接，支持连接、断开、发送个人消息和广播消息功能。
    """
    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    @staticmethod
    async def send_personal_message(message, websocket: WebSocket):
        """发送个人消息，统一接收Python对象"""
        if isinstance(message, str):
            # 如果传入字符串，尝试解析为JSON对象
            try:
                message = json.loads(message)
            except json.JSONDecodeError:
                # 解析失败，发送文本消息
                await websocket.send_text(message)
                return
        await websocket.send_json(message)

    async def broadcast(self, message, exclude_websocket: WebSocket = None, *args, **kwargs):
        """广播消息给所有连接的客户端，统一接收Python对象"""
        disconnected_connections = []
        for connection in self.active_connections:
            if exclude_websocket and connection == exclude_websocket:
                continue
            try:
                if isinstance(message, str):
                    # 如果传入字符串，尝试解析为JSON对象
                    try:
                        message_data = json.loads(message)
                        await connection.send_json(message_data)
                    except json.JSONDecodeError:
                        # 解析失败，发送文本消息
                        await connection.send_text(message)
                else:
                    # 直接发送Python对象
                    await connection.send_json(message)
            except Exception as e:
                # 连接已断开，记录下来稍后移除
                disconnected_connections.append(connection)
        
        # 移除已断开的连接
        for connection in disconnected_connections:
            if connection in self.active_connections:
                self.active_connections.remove(connection)
